package com.dxy.elasticsearch.config;

import lombok.extern.slf4j.Slf4j;
import org.apache.http.HttpHost;
import org.apache.http.HttpResponse;
import org.apache.http.auth.AuthScope;
import org.apache.http.auth.Credentials;
import org.apache.http.auth.UsernamePasswordCredentials;
import org.apache.http.client.CredentialsProvider;
import org.apache.http.impl.client.BasicCredentialsProvider;
import org.apache.http.impl.client.DefaultConnectionKeepAliveStrategy;
import org.apache.http.protocol.HttpContext;
import org.elasticsearch.client.RestClient;
import org.elasticsearch.client.RestClientBuilder;
import org.elasticsearch.client.RestHighLevelClient;
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.util.StringUtils;

import java.util.Arrays;
import java.util.Objects;

/**
 * @author wudapei
 * @date 2022-11-21 14:03
 */
@Slf4j
@Configuration
@EnableConfigurationProperties({EsProperties.class})
public class ElasticsearchConfig {

    @Bean
    @ConditionalOnMissingBean({RestHighLevelClient.class})
    public RestHighLevelClient restHighLevelClient(EsProperties esProperties) {
        RestClientBuilder builder = RestClient.builder(Objects.requireNonNull(this.parseEsNodes(esProperties.getNodes())));
        builder.setRequestConfigCallback((requestConfigBuilder) -> {
            requestConfigBuilder.setConnectTimeout(esProperties.getConnectTimeoutMs());
            requestConfigBuilder.setSocketTimeout(esProperties.getSocketTimeoutMs());
            requestConfigBuilder.setConnectionRequestTimeout(esProperties.getConnectionRequestTimeoutMs());
            return requestConfigBuilder;
        });
        builder.setHttpClientConfigCallback((httpClientBuilder) -> {
            if (StringUtils.hasText(esProperties.getUsername())) {
                CredentialsProvider provider = new BasicCredentialsProvider();
                Credentials credentials = new UsernamePasswordCredentials(esProperties.getUsername(), esProperties.getPassword());
                provider.setCredentials(AuthScope.ANY, credentials);
                httpClientBuilder.setDefaultCredentialsProvider(provider);
            }

            httpClientBuilder.setMaxConnTotal(esProperties.getMaxConnTotal());
            httpClientBuilder.setMaxConnPerRoute(esProperties.getMaxConnPerRoute());
            httpClientBuilder.setKeepAliveStrategy(new DefaultConnectionKeepAliveStrategy() {
                @Override
                public long getKeepAliveDuration(final HttpResponse response, final HttpContext context) {
                    long keepAlive = super.getKeepAliveDuration(response, context);
                    if (keepAlive == -1L) {
                        keepAlive = esProperties.getKeepAliveMs();
                    }

                    return keepAlive;
                }
            });
            return httpClientBuilder;
        });
        log.info("初始化Elasticsearch RestHighLevelClient");
        return new RestHighLevelClient(builder);
    }

    private HttpHost[] parseEsNodes(String nodes) {
        log.info("Elasticsearch服务器地址配置:{}", nodes);
        if (org.apache.commons.lang3.StringUtils.isBlank(nodes)) {
            throw new IllegalArgumentException("Elasticsearch服务器地址为空，请检查配置是否正确！");
        }
        return Arrays.stream(nodes.split(",")).filter(org.apache.commons.lang3.StringUtils::isNotBlank).map(HttpHost::create).toArray(HttpHost[]::new);
    }
}
